#!/usr/bin/python3
# -*- coding: utf-8 -*-

import sys
# sys.path.append('/opt/work/caffe/python')
sys.path.insert(0, '.')

import os
import random


file_list_path = r'/rootfs/media/kasim/Data1/data/VideoGenderCropPerson/faceid_video_gender_age.txt'
file_gender_list_path = r'/rootfs/media/kasim/Data1/data/VideoGenderCropPerson/person_image_gender.txt'


val_ratio = 0.15
gender_ratio = 4/1

UNDER_TWO = 0                       # 2<=
THREE_TO_FIVE = 1                   # 3-5
SIX_TO_TWELVE = 2                   # 6-12
THIRTEEN_TO_EIGHTEEN = 3            # 13-18
NINETEEN_TO_TWENTY_THREE = 4        # 19-23
TWENTY_FOUR_TO_TWENTY_NINE = 5      # 24-29
THIRTY_TO_THIRTY_FIVE = 6           # 30-35
THIRTY_SIX_TO_FORTY_FIVE = 7        # 36-45
FORTY_SIX_TO_SIXTY = 8              # 46-60
ABOVE_SIXTY = 9                     # 60>
UNCERTAIN = 10


AGE_1_10 = 0
AGE_11_18 = 1
AGE_19_25 = 2
AGE_26_35 = 3
AGE_35_45 = 4
AGE_45_55 = 5
AGE_56_65 = 6
AGE_66 = 7


def main():
    random.seed(123456)

    gender_age_group_faceid = {}
    age_group_gender_faceid = {}
    faceid_video_gender_age_agegroups = {}
    with open(file_list_path, 'r') as file:
        for line in file.readlines():
            # 1007 1007 gd/zs/dqjd/xgly/150100414a54443452067fa6d4ad6600/video/20190925_230722,gd/zs/dqjd/xgly/150100414a54443452067fa6d4ad6600/video/20190925_230919 1 62.26731872558594 9
            lines = line.strip().split()
            faceid = '{} {}'.format(lines[0], lines[1])
            videos = lines[2].split(',')
            gender = int(lines[3])
            age = float(lines[4])
            age_group = int(lines[5])
            faceid_video_gender_age_agegroups[faceid] = [videos, gender, age, age_group]

            age_group_faceids = gender_age_group_faceid.get(gender, None)
            if age_group_faceids is None:
                gender_age_group_faceid[gender] = {age_group: [faceid]}
            else:
                faceids = age_group_faceids.get(age_group, None)
                if faceids is None:
                    age_group_faceids[age_group] = [faceid]
                else:
                    faceids.append(faceid)

            gender_faceids = age_group_gender_faceid.get(age_group, None)
            if gender_faceids is None:
                age_group_gender_faceid[age_group] = {gender: [faceid]}
            else:
                faceids = gender_faceids.get(gender, None)
                if faceids is None:
                    gender_faceids[gender] = [faceid]
                else:
                    faceids.append(faceid)

    # UNDER_TWO = 0  # 2<=
    # THREE_TO_FIVE = 1  # 3-5
    # SIX_TO_TWELVE = 2  # 6-12
    # THIRTEEN_TO_EIGHTEEN = 3  # 13-18
    # NINETEEN_TO_TWENTY_THREE = 4  # 19-23
    # TWENTY_FOUR_TO_TWENTY_NINE = 5  # 24-29
    # THIRTY_TO_THIRTY_FIVE = 6  # 30-35
    # THIRTY_SIX_TO_FORTY_FIVE = 7  # 36-45
    # FORTY_SIX_TO_SIXTY = 8  # 46-60
    # ABOVE_SIXTY = 9  # 60>
    # UNCERTAIN = 10

    print('*' * 80)
    format_str = '{} - {}~{}Age: {}'
    for gender, age_group_faceid in gender_age_group_faceid.items():
        gender = 'FeMale' if gender == 0 else 'Male'
        for age_group, faceids in age_group_faceid.items():
            faceid_count = len(faceids)

            if UNDER_TWO == age_group:
                print(format_str.format(gender, 0, 2, faceid_count))
            elif THREE_TO_FIVE == age_group:
                print(format_str.format(gender, 3, 5, faceid_count))
            elif SIX_TO_TWELVE == age_group:
                print(format_str.format(gender, 6, 12, faceid_count))
            elif THIRTEEN_TO_EIGHTEEN == age_group:
                print(format_str.format(gender, 13, 18, faceid_count))
            elif NINETEEN_TO_TWENTY_THREE == age_group:
                print(format_str.format(gender, 19, 23, faceid_count))
            elif TWENTY_FOUR_TO_TWENTY_NINE == age_group:
                print(format_str.format(gender, 24, 29, faceid_count))
            elif THIRTY_TO_THIRTY_FIVE == age_group:
                print(format_str.format(gender, 30, 35, faceid_count))
            elif THIRTY_SIX_TO_FORTY_FIVE == age_group:
                print(format_str.format(gender, 36, 45, faceid_count))
            elif FORTY_SIX_TO_SIXTY == age_group:
                print(format_str.format(gender, 46, 60, faceid_count))
            elif ABOVE_SIXTY == age_group:
                print(format_str.format(gender, 60, 99, faceid_count))
            elif UNCERTAIN == age_group:
                print(format_str.format(gender, -1, -1, faceid_count))

    print('*' * 80)
    format_str = '{}~{}Age - {}: {}'
    for age_group, gender_faceid in age_group_gender_faceid.items():
        for gender, faceids in gender_faceid.items():
            faceid_count = len(faceids)

            gender = 'FeMale' if gender == 0 else 'Male'
            if UNDER_TWO == age_group:
                print(format_str.format(0, 2, gender, faceid_count))
            elif THREE_TO_FIVE == age_group:
                print(format_str.format(3, 5, gender, faceid_count))
            elif SIX_TO_TWELVE == age_group:
                print(format_str.format(6, 12, gender, faceid_count))
            elif THIRTEEN_TO_EIGHTEEN == age_group:
                print(format_str.format(13, 18, gender, faceid_count))
            elif NINETEEN_TO_TWENTY_THREE == age_group:
                print(format_str.format(19, 23, gender, faceid_count))
            elif TWENTY_FOUR_TO_TWENTY_NINE == age_group:
                print(format_str.format(24, 29, gender, faceid_count))
            elif THIRTY_TO_THIRTY_FIVE == age_group:
                print(format_str.format(30, 35, gender, faceid_count))
            elif THIRTY_SIX_TO_FORTY_FIVE == age_group:
                print(format_str.format(36, 45, gender, faceid_count))
            elif FORTY_SIX_TO_SIXTY == age_group:
                print(format_str.format(46, 60, gender, faceid_count))
            elif ABOVE_SIXTY == age_group:
                print(format_str.format(60, 99, gender, faceid_count))
            elif UNCERTAIN == age_group:
                print(format_str.format(-1, -1, gender, faceid_count))

    print('*' * 80)

    out_path = os.path.dirname(file_list_path)

    train_file_path = os.path.join(out_path, 'train_faceid_video_gender_age.txt')
    val_file_path = os.path.join(out_path, 'val_faceid_video_gender_age.txt')

    train_file = open(train_file_path, 'w')
    val_file = open(val_file_path, 'w')

    train_videos = set()
    val_videos = set()
    format_str = '{}~{}Age: Train Female {}, Train Male {}, Val Female {}, Val Male {}'
    for age_group, gender_faceid in age_group_gender_faceid.items():
        female_faceids = gender_faceid.get(0, None)
        male_faceids = gender_faceid.get(1, None)
        if (female_faceids is None) or (male_faceids is None):
            continue
        random.shuffle(female_faceids)
        random.shuffle(male_faceids)

        female_faceid_count = len(female_faceids)
        male_faceid_count = len(male_faceids)

        val_female_count = int(round(val_ratio * female_faceid_count))
        if val_female_count < 1:
            val_female_count = 1
        train_female_count = female_faceid_count - val_female_count
        if (val_female_count == 0) and (train_female_count == 0):
            continue

        val_male_count = int(round(val_ratio * male_faceid_count))
        if val_male_count < 1:
            val_male_count = 1
        train_male_count = male_faceid_count - val_male_count
        if (val_male_count == 0) and (train_male_count == 0):
            continue

        if val_female_count > val_male_count:
            max_val_female_count = int(round(gender_ratio * val_male_count))
            if val_female_count > max_val_female_count:
                val_female_count = max_val_female_count
        elif val_female_count < val_male_count:
            max_val_male_count = int(round(gender_ratio * val_female_count))
            if val_male_count > max_val_male_count:
                val_male_count = max_val_male_count

        if train_female_count > train_male_count:
            max_train_female_count = int(round(gender_ratio * train_male_count))
            if train_female_count > max_train_female_count:
                train_female_count = max_train_female_count
        elif train_female_count < train_male_count:
            max_train_male_count = int(round(gender_ratio * train_female_count))
            if train_male_count > max_train_male_count:
                train_male_count = max_train_male_count

        train_female_faceids = female_faceids[0:train_female_count]
        val_female_faceids = female_faceids[train_female_count:(train_female_count+val_female_count)]
        train_male_faceids = male_faceids[0:train_male_count]
        val_male_faceids = male_faceids[train_male_count:(train_male_count+val_male_count)]

        train_female_count = len(train_female_faceids)
        val_female_count = len(val_female_faceids)
        train_male_count = len(train_male_faceids)
        val_male_count = len(val_male_faceids)

        if UNDER_TWO == age_group:
            print(format_str.format(0, 2, train_female_count, train_male_count, val_female_count, val_male_count))
        elif THREE_TO_FIVE == age_group:
            print(format_str.format(3, 5, train_female_count, train_male_count, val_female_count, val_male_count))
        elif SIX_TO_TWELVE == age_group:
            print(format_str.format(6, 12, train_female_count, train_male_count, val_female_count, val_male_count))
        elif THIRTEEN_TO_EIGHTEEN == age_group:
            print(format_str.format(13, 18, train_female_count, train_male_count, val_female_count, val_male_count))
        elif NINETEEN_TO_TWENTY_THREE == age_group:
            print(format_str.format(19, 23, train_female_count, train_male_count, val_female_count, val_male_count))
        elif TWENTY_FOUR_TO_TWENTY_NINE == age_group:
            print(format_str.format(24, 29, train_female_count, train_male_count, val_female_count, val_male_count))
        elif THIRTY_TO_THIRTY_FIVE == age_group:
            print(format_str.format(30, 35, train_female_count, train_male_count, val_female_count, val_male_count))
        elif THIRTY_SIX_TO_FORTY_FIVE == age_group:
            print(format_str.format(36, 45, train_female_count, train_male_count, val_female_count, val_male_count))
        elif FORTY_SIX_TO_SIXTY == age_group:
            print(format_str.format(46, 60, train_female_count, train_male_count, val_female_count, val_male_count))
        elif ABOVE_SIXTY == age_group:
            print(format_str.format(60, 99, train_female_count, train_male_count, val_female_count, val_male_count))
        elif UNCERTAIN == age_group:
            print(format_str.format(-1, -1, train_female_count, train_male_count, val_female_count, val_male_count))

        for faceid in train_female_faceids:
            video_gender_age_agegroups = faceid_video_gender_age_agegroups[faceid]
            videos, gender, age, agegroups = video_gender_age_agegroups
            train_videos.update(videos)
            videos_str = ','.join(videos)
            out_info = '{} {} {} {} {}\n'.format(faceid, videos_str, gender, age, agegroups)
            train_file.write(out_info)

        for faceid in train_male_faceids:
            video_gender_age_agegroups = faceid_video_gender_age_agegroups[faceid]
            videos, gender, age, agegroups = video_gender_age_agegroups
            train_videos.update(videos)
            videos_str = ','.join(videos)
            out_info = '{} {} {} {} {}\n'.format(faceid, videos_str, gender, age, agegroups)
            train_file.write(out_info)

        for faceid in val_female_faceids:
            video_gender_age_agegroups = faceid_video_gender_age_agegroups[faceid]
            videos, gender, age, agegroups = video_gender_age_agegroups
            val_videos.update(videos)
            videos_str = ','.join(videos)
            out_info = '{} {} {} {} {}\n'.format(faceid, videos_str, gender, age, agegroups)
            val_file.write(out_info)

        for faceid in val_male_faceids:
            video_gender_age_agegroups = faceid_video_gender_age_agegroups[faceid]
            videos, gender, age, agegroups = video_gender_age_agegroups
            val_videos.update(videos)
            videos_str = ','.join(videos)
            out_info = '{} {} {} {} {}\n'.format(faceid, videos_str, gender, age, agegroups)
            val_file.write(out_info)

    train_file.close()
    val_file.close()
    os.system('chmod a+wr {}'.format(train_file_path))
    os.system('chmod a+wr {}'.format(val_file_path))

    train_image_genders = []
    val_image_genders = []
    with open(file_gender_list_path, 'r') as file:
        for line in file.readlines():
            video_name = os.path.dirname(line.split()[0])
            if video_name in train_videos:
                train_image_genders.append(line)
            elif video_name in val_videos:
                val_image_genders.append(line)

    train_file_path = os.path.join(out_path, 'train_person_image_gender.txt')
    val_file_path = os.path.join(out_path, 'val_person_image_gender.txt')

    train_file = open(train_file_path, 'w')
    val_file = open(val_file_path, 'w')

    random.shuffle(train_image_genders)
    random.shuffle(val_image_genders)

    for line in train_image_genders:
        train_file.write(line)

    for line in val_image_genders:
        val_file.write(line)

    train_file.close()
    val_file.close()
    os.system('chmod a+wr {}'.format(train_file_path))
    os.system('chmod a+wr {}'.format(val_file_path))

    print('Train Video Count: {}, Train Image Count: {}, Val Video Count: {}, Val Image Count: {}'.format(len(train_videos), len(train_image_genders), len(val_videos), len(val_image_genders)))

    print('Finish!')


if __name__ == '__main__':
    main()
